{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Tce3stUlHN0L"
   },
   "source": [
    "##### Copyright 2024 Google LLC."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "cellView": "form",
    "id": "tuOe1ymfHZPu"
   },
   "outputs": [],
   "source": [
    "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PXNm5_p_oxMF"
   },
   "source": [
    "# Gemma - Run with Mesop\n",
    "\n",
    "This notebook demonstrates how you can run a Gemma model with [Google Mesop](https://github.com/google/mesop). Mesop is a Python-based UI framework that allows you to rapidly build web apps like demos and internal apps.\n",
    "\n",
    "<table align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/Gemma/Integrate_with_Mesop.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mERVCCsGUPIJ"
   },
   "source": [
    "## Setup\n",
    "\n",
    "### Select the Colab runtime\n",
    "To complete this tutorial, you'll need to have a Colab runtime with sufficient resources to run the Gemma model. In this case, you can use a T4 GPU:\n",
    "\n",
    "1. In the upper-right of the Colab window, select **▾ (Additional connection options)**.\n",
    "2. Select **Change runtime type**.\n",
    "3. Under **Hardware accelerator**, select **T4 GPU**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QQ6W7NzRe1VM"
   },
   "source": [
    "### Gemma setup\n",
    "\n",
    "To complete this tutorial, you'll first need to complete the setup instructions at [Gemma setup](https://ai.google.dev/gemma/docs/setup). The Gemma setup instructions show you how to do the following:\n",
    "\n",
    "* Get access to Gemma on kaggle.com.\n",
    "* Select a Colab runtime with sufficient resources to run\n",
    "  the Gemma 2B model.\n",
    "* Generate and configure a Kaggle username and API key.\n",
    "\n",
    "After you've completed the Gemma setup, move on to the next section, where you'll set environment variables for your Colab environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_gN-IVRC3dQe"
   },
   "source": [
    "### Set environment variables\n",
    "\n",
    "Set environment variables for `KAGGLE_USERNAME` and `KAGGLE_KEY`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "DrBoa_Urw9Vx"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from google.colab import userdata\n",
    "\n",
    "# Note: `userdata.get` is a Colab API. If you're not using Colab, set the env\n",
    "# vars as appropriate for your system.\n",
    "os.environ[\"KAGGLE_USERNAME\"] = userdata.get(\"KAGGLE_USERNAME\")\n",
    "os.environ[\"KAGGLE_KEY\"] = userdata.get(\"KAGGLE_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z9oy3QUmXtSd"
   },
   "source": [
    "### Install dependencies\n",
    "\n",
    "You will run KerasNLP to run Gemma. So install Keras and KerasNLP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "UcGLzDeQ8NwN",
    "outputId": "e8990335-db7d-4852-caeb-a2810a4a7671",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m570.5/570.5 kB\u001b[0m \u001b[31m7.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m950.8/950.8 kB\u001b[0m \u001b[31m12.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.2/5.2 MB\u001b[0m \u001b[31m28.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m589.8/589.8 MB\u001b[0m \u001b[31m2.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.3/5.3 MB\u001b[0m \u001b[31m107.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m100.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.5/5.5 MB\u001b[0m \u001b[31m114.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m76.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m311.2/311.2 kB\u001b[0m \u001b[31m37.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "tf-keras 2.15.1 requires tensorflow<2.16,>=2.15, but you have tensorflow 2.16.1 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "# Install Keras 3 last. See https://keras.io/getting_started/ for more details.\n",
    "!pip install -q -U keras-nlp\n",
    "!pip install -q -U keras>=3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Pm5cVOFt5YvZ"
   },
   "source": [
    "### Select a backend\n",
    "\n",
    "You will use the JAX backend for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "7rS7ryTs5wjf"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"KERAS_BACKEND\"] = \"jax\"  # Or \"tensorflow\" or \"torch\".\n",
    "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"0.9\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "599765c72722"
   },
   "source": [
    "### Import packages\n",
    "\n",
    "Import Keras and KerasNLP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "f2fa267d75bc"
   },
   "outputs": [],
   "source": [
    "import keras\n",
    "import keras_nlp"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Enable mixed precision on GPU."
   ],
   "metadata": {
    "id": "xfUlIT24giK8"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "keras.mixed_precision.set_global_policy(\"mixed_bfloat16\")"
   ],
   "metadata": {
    "id": "s79GrIXQf2HS"
   },
   "execution_count": 6,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZsxDCbLN555T"
   },
   "source": [
    "## Create a model\n",
    "\n",
    "Create the Gemma model using the `from_preset` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "yygIK9DEIldp",
    "outputId": "105344d1-95ee-46ca-f631-e68fb99ed535",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'task.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'config.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'config.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'config.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'model.weights.h5' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'preprocessor.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n",
      "Attaching 'assets/tokenizer/vocabulary.spm' from model 'keras/gemma/keras/gemma_1.1_instruct_2b_en/3' to your Colab notebook...\n"
     ]
    }
   ],
   "source": [
    "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(\"gemma_1.1_instruct_2b_en\")"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Install and start Mesop"
   ],
   "metadata": {
    "id": "y5EMEQJgnfus"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "!pip install mesop\n",
    "import mesop as me\n",
    "import mesop.labs as mel\n",
    "\n",
    "me.colab_run()"
   ],
   "metadata": {
    "id": "KK-WfjB11DqO",
    "outputId": "5f276aca-815d-4838-a389-2ad6daaf2611",
    "colab": {
     "base_uri": "https://localhost:8080/"
    }
   },
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Collecting mesop\n",
      "  Downloading mesop-0.8.0-py3-none-any.whl (5.0 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.0/5.0 MB\u001b[0m \u001b[31m17.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: absl-py in /usr/local/lib/python3.10/dist-packages (from mesop) (1.4.0)\n",
      "Collecting deepdiff==6.7.1 (from mesop)\n",
      "  Downloading deepdiff-6.7.1-py3-none-any.whl (76 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m76.6/76.6 kB\u001b[0m \u001b[31m13.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: flask in /usr/local/lib/python3.10/dist-packages (from mesop) (2.2.5)\n",
      "Collecting libcst==1.1.0 (from mesop)\n",
      "  Downloading libcst-1.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.2 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.2/3.2 MB\u001b[0m \u001b[31m42.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: markdown in /usr/local/lib/python3.10/dist-packages (from mesop) (3.6)\n",
      "Requirement already satisfied: protobuf in /usr/local/lib/python3.10/dist-packages (from mesop) (3.20.3)\n",
      "Collecting pydantic==1.10.13 (from mesop)\n",
      "  Downloading pydantic-1.10.13-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.1 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.1/3.1 MB\u001b[0m \u001b[31m58.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: pygments in /usr/local/lib/python3.10/dist-packages (from mesop) (2.16.1)\n",
      "Collecting watchdog (from mesop)\n",
      "  Downloading watchdog-4.0.1-py3-none-manylinux2014_x86_64.whl (83 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m83.0/83.0 kB\u001b[0m \u001b[31m14.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting ordered-set<4.2.0,>=4.0.2 (from deepdiff==6.7.1->mesop)\n",
      "  Downloading ordered_set-4.1.0-py3-none-any.whl (7.6 kB)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.2 in /usr/local/lib/python3.10/dist-packages (from libcst==1.1.0->mesop) (4.12.1)\n",
      "Collecting typing-inspect>=0.4.0 (from libcst==1.1.0->mesop)\n",
      "  Downloading typing_inspect-0.9.0-py3-none-any.whl (8.8 kB)\n",
      "Requirement already satisfied: pyyaml>=5.2 in /usr/local/lib/python3.10/dist-packages (from libcst==1.1.0->mesop) (6.0.1)\n",
      "Requirement already satisfied: Werkzeug>=2.2.2 in /usr/local/lib/python3.10/dist-packages (from flask->mesop) (3.0.3)\n",
      "Requirement already satisfied: Jinja2>=3.0 in /usr/local/lib/python3.10/dist-packages (from flask->mesop) (3.1.4)\n",
      "Requirement already satisfied: itsdangerous>=2.0 in /usr/local/lib/python3.10/dist-packages (from flask->mesop) (2.2.0)\n",
      "Requirement already satisfied: click>=8.0 in /usr/local/lib/python3.10/dist-packages (from flask->mesop) (8.1.7)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from Jinja2>=3.0->flask->mesop) (2.1.5)\n",
      "Collecting mypy-extensions>=0.3.0 (from typing-inspect>=0.4.0->libcst==1.1.0->mesop)\n",
      "  Downloading mypy_extensions-1.0.0-py3-none-any.whl (4.7 kB)\n",
      "Installing collected packages: watchdog, pydantic, ordered-set, mypy-extensions, typing-inspect, deepdiff, libcst, mesop\n",
      "  Attempting uninstall: pydantic\n",
      "    Found existing installation: pydantic 2.7.3\n",
      "    Uninstalling pydantic-2.7.3:\n",
      "      Successfully uninstalled pydantic-2.7.3\n",
      "Successfully installed deepdiff-6.7.1 libcst-1.1.0 mesop-0.8.0 mypy-extensions-1.0.0 ordered-set-4.1.0 pydantic-1.10.13 typing-inspect-0.9.0 watchdog-4.0.1\n",
      "\n",
      "\u001b[32mRunning server on: http://localhost:32123\u001b[0m\n",
      " * Serving Flask app 'mesop.server.server'\n",
      " * Debug mode: off\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Load the Mesop UI."
   ],
   "metadata": {
    "id": "0C2CnGxlnlsK"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "@me.page(path=\"/chat\")\n",
    "def chat():\n",
    "    mel.chat(transform)\n",
    "\n",
    "\n",
    "def transform(user_prompt: str, history: list[mel.ChatMessage]) -> str:\n",
    "\n",
    "    # Assemble prompt from chat history\n",
    "    prompt = \"\"\n",
    "    for h in history:\n",
    "        prompt += \"<start_of_turn>{role}\\n{content}<end_of_turn>\\n\".format(\n",
    "            role=h.role, content=h.content\n",
    "        )\n",
    "    prompt += \"<start_of_turn>model\\n\"\n",
    "\n",
    "    result = gemma_lm.generate(prompt)\n",
    "    return result[len(prompt) :]\n",
    "\n",
    "\n",
    "me.colab_show(path=\"/chat\")"
   ],
   "metadata": {
    "id": "D1hVvC5b1KI-",
    "outputId": "50a8a149-83aa-4c31-d553-26d4a5c40642",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 510
    }
   },
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ],
      "application/javascript": [
       "(async (port, path, width, height, cache, element) => {\n",
       "    if (!google.colab.kernel.accessAllowed && !cache) {\n",
       "      return;\n",
       "    }\n",
       "    element.appendChild(document.createTextNode(''));\n",
       "    const url = await google.colab.kernel.proxyPort(port, {cache});\n",
       "    const iframe = document.createElement('iframe');\n",
       "    iframe.src = new URL(path, url).toString();\n",
       "    iframe.height = height;\n",
       "    iframe.width = width;\n",
       "    iframe.style.border = 0;\n",
       "    iframe.allow = [\n",
       "        'accelerometer',\n",
       "        'autoplay',\n",
       "        'camera',\n",
       "        'clipboard-read',\n",
       "        'clipboard-write',\n",
       "        'gyroscope',\n",
       "        'magnetometer',\n",
       "        'microphone',\n",
       "        'serial',\n",
       "        'usb',\n",
       "        'xr-spatial-tracking',\n",
       "    ].join('; ');\n",
       "    element.appendChild(iframe);\n",
       "  })(32123, \"/chat\", \"100%\", \"400\", false, window.element)"
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "INFO:werkzeug:\u001b[31m\u001b[1mWARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.\u001b[0m\n",
      " * Running on all addresses (::)\n",
      " * Running on http://[::1]:32123\n",
      " * Running on http://[::1]:32123\n",
      "INFO:werkzeug:\u001b[33mPress CTRL+C to quit\u001b[0m\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now you can chat with the Gemma model in the Mesop UI. You can restart the conversation by running the cell above again."
   ],
   "metadata": {
    "id": "J7enAHFb1C4i"
   }
  },
  {
   "cell_type": "code",
   "source": [],
   "metadata": {
    "id": "WAWz28QLk24r"
   },
   "execution_count": 9,
   "outputs": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "machine_shape": "hm",
   "gpuType": "T4"
  },
  "google": {
   "image_path": "/site-assets/images/marketing/gemma.png",
   "keywords": [
    "examples",
    "gemma",
    "python",
    "quickstart",
    "text"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "accelerator": "GPU"
 },
 "nbformat": 4,
 "nbformat_minor": 0
}